#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import argparse
import torch
import numpy as np
from tqdm import tqdm
import time
from pathlib import Path
import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import threading
import queue
import signal
import keyboard

# 设置matplotlib支持中文显示
import matplotlib
matplotlib.rcParams['font.family'] = ['SimHei', 'Microsoft YaHei', 'sans-serif']  # 优先使用黑体、微软雅黑
matplotlib.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
# 如果系统中没有中文字体，可以尝试使用系统默认字体
try:
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
except:
    print("警告: 无法设置中文字体，可能导致中文显示为乱码")

# 导入模型定义
from transformer_denoising_model import PointCloudTransformerDenoiser
from denoise_pointcloud import PointCloudDenoiser

class RealTimePointCloudDenoiser:
    def __init__(self, model_path, device='auto', buffer_size=5, visualization=True):
        """
        实时点云去噪器
        
        Args:
            model_path: 预训练模型的路径
            device: 运行设备 ('cuda', 'cpu' 或 'auto')
            buffer_size: 处理缓冲区大小
            visualization: 是否启用实时可视化
        """
        # 创建去噪器
        self.denoiser = PointCloudDenoiser(model_path, device)
        self.device = self.denoiser.device
        self.model = self.denoiser.model
        
        # 设置缓冲区
        self.buffer_size = buffer_size
        self.point_cloud_queue = queue.Queue(maxsize=buffer_size)
        
        # 处理状态
        self.running = False
        self.total_frames = 0
        self.total_time = 0
        self.visualization = visualization
        
        # 如果启用可视化，设置matplotlib
        if visualization:
            plt.ion()  # 开启交互模式
            self.fig = plt.figure(figsize=(15, 5))
            
            # 原始点云
            self.ax1 = self.fig.add_subplot(121, projection='3d')
            self.ax1.set_title('原始点云')
            self.ax1.set_xlabel('X')
            self.ax1.set_ylabel('Y')
            self.ax1.set_zlabel('Z')
            
            # 去噪后点云
            self.ax2 = self.fig.add_subplot(122, projection='3d')
            self.ax2.set_title('去噪点云')
            self.ax2.set_xlabel('X')
            self.ax2.set_ylabel('Y')
            self.ax2.set_zlabel('Z')
            
            plt.tight_layout()
    
    def preprocess_lidar_data(self, points):
        """
        预处理LiDAR数据
        
        Args:
            points: 原始点云数据 [N, 3+]
            
        Returns:
            预处理后的点云张量
        """
        # 提取XYZ坐标
        if points.shape[1] > 3:
            points = points[:, :3]
        
        # 归一化点云
        center = np.mean(points, axis=0)
        points = points - center
        
        scale = np.max(np.linalg.norm(points, axis=1))
        if scale > 0:
            normalized_points = points / scale
        else:
            normalized_points = points
        
        # 确保点的数量是2048
        target_num_points = 2048
        
        if len(normalized_points) > target_num_points:
            # 随机采样
            indices = np.random.choice(len(normalized_points), target_num_points, replace=False)
            sampled_points = normalized_points[indices]
        elif len(normalized_points) < target_num_points:
            # 重复点以达到目标数量
            repeat_factor = int(np.ceil(target_num_points / len(normalized_points)))
            repeated_points = np.repeat(normalized_points, repeat_factor, axis=0)
            sampled_points = repeated_points[:target_num_points]
        else:
            sampled_points = normalized_points
        
        # 转换为tensor
        return torch.tensor(sampled_points, dtype=torch.float32), center, scale
    
    def postprocess_points(self, denoised_points, center, scale, additional_attributes=None):
        """
        后处理去噪后的点云
        
        Args:
            denoised_points: 去噪后的点云张量
            center: 原始点云的中心
            scale: 原始点云的缩放因子
            additional_attributes: 点云的其他属性 (颜色、强度等)
            
        Returns:
            后处理的点云
        """
        # 转换为numpy数组
        points = denoised_points.cpu().numpy()
        
        # 反归一化
        if scale > 0:
            points = points * scale + center
        else:
            points = points + center
        
        # 如果有其他属性，添加回去
        if additional_attributes is not None:
            # 这里需要根据实际情况处理,例如基于最近邻匹配
            pass
        
        return points
    
    def denoise_frame(self, point_cloud):
        """
        对单帧点云进行去噪
        
        Args:
            point_cloud: 输入的点云帧
            
        Returns:
            去噪后的点云
        """
        # 预处理点云
        normalized_points, center, scale = self.preprocess_lidar_data(point_cloud)
        
        # 添加批次维度
        batch_points = normalized_points.unsqueeze(0).to(self.device)
        
        # 计时
        start_time = time.time()
        
        # 去噪
        with torch.no_grad():
            denoised_batch_points = self.model(batch_points)
        
        # 计算处理时间
        elapsed = time.time() - start_time
        
        # 后处理
        denoised_points = self.postprocess_points(denoised_batch_points[0], center, scale)
        
        return denoised_points, elapsed
    
    def visualize_frame(self, original_points, denoised_points):
        """
        可视化点云帧
        
        Args:
            original_points: 原始点云
            denoised_points: 去噪后的点云
        """
        # 清除当前轴
        self.ax1.clear()
        self.ax2.clear()
        
        # 设置标题和坐标轴标签
        self.ax1.set_title('原始点云')
        self.ax1.set_xlabel('X')
        self.ax1.set_ylabel('Y')
        self.ax1.set_zlabel('Z')
        
        self.ax2.set_title('去噪点云')
        self.ax2.set_xlabel('X')
        self.ax2.set_ylabel('Y')
        self.ax2.set_zlabel('Z')
        
        # 绘制点云
        self.ax1.scatter(original_points[:, 0], original_points[:, 1], original_points[:, 2], 
                        c=original_points[:, 2], cmap='viridis', s=1)
        self.ax2.scatter(denoised_points[:, 0], denoised_points[:, 1], denoised_points[:, 2], 
                        c=denoised_points[:, 2], cmap='plasma', s=1)
        
        # 设置相同的视角
        self.ax2.view_init(elev=self.ax1.elev, azim=self.ax1.azim)
        
        # 刷新图形
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()
    
    def process_thread(self):
        """处理线程函数"""
        print("启动处理线程...")
        
        while self.running:
            try:
                # 从队列获取点云帧，设置超时以便能够响应结束信号
                try:
                    frame_data = self.point_cloud_queue.get(timeout=1.0)
                except queue.Empty:
                    continue
                
                # 处理点云
                original_points = frame_data["points"]
                denoised_points, elapsed = self.denoise_frame(original_points)
                
                # 更新统计信息
                self.total_frames += 1
                self.total_time += elapsed
                avg_time = self.total_time / self.total_frames
                fps = 1.0 / avg_time if avg_time > 0 else 0
                
                print(f"\r处理第 {self.total_frames} 帧 | 当前帧处理时间: {elapsed:.4f}秒 | 平均处理时间: {avg_time:.4f}秒 | FPS: {fps:.2f}", end="")
                
                # 如果启用可视化，显示当前帧
                if self.visualization:
                    self.visualize_frame(original_points, denoised_points)
                
                # 如果有回调函数，调用它
                if frame_data.get("callback"):
                    frame_data["callback"](denoised_points)
                
                # 标记为完成
                self.point_cloud_queue.task_done()
                
            except Exception as e:
                print(f"\n处理线程发生错误: {e}")
                continue
        
        print("\n处理线程结束")
    
    def start(self):
        """启动实时去噪处理"""
        if self.running:
            print("处理器已经在运行中")
            return
        
        self.running = True
        self.total_frames = 0
        self.total_time = 0
        
        # 创建处理线程
        self.process_thread_handle = threading.Thread(target=self.process_thread)
        self.process_thread_handle.daemon = True
        self.process_thread_handle.start()
        
        print("实时点云去噪器已启动")
    
    def stop(self):
        """停止实时处理"""
        self.running = False
        
        # 等待处理线程结束
        if hasattr(self, 'process_thread_handle') and self.process_thread_handle.is_alive():
            self.process_thread_handle.join(timeout=2.0)
        
        print("\n实时点云去噪器已停止")
        
        # 关闭可视化
        if self.visualization:
            plt.close(self.fig)
            plt.ioff()
    
    def add_point_cloud(self, points, callback=None):
        """
        添加点云到处理队列
        
        Args:
            points: 点云数据，numpy数组 [N, 3+]
            callback: 处理完成后的回调函数
            
        Returns:
            是否成功添加到队列
        """
        if not self.running:
            print("处理器未启动")
            return False
        
        if self.point_cloud_queue.full():
            print("队列已满，跳过当前帧")
            return False
        
        # 添加到队列
        self.point_cloud_queue.put({
            "points": points,
            "callback": callback
        })
        
        return True


def simulate_lidar_data(num_points=1000):
    """
    模拟LiDAR数据
    
    Args:
        num_points: 模拟的点数量
        
    Returns:
        模拟的点云数据
    """
    # 模拟一个旋转的圆环点云
    theta = np.random.uniform(0, 2*np.pi, num_points)
    r = np.random.normal(5.0, 0.2, num_points)  # 半径围绕5米，有些噪声
    
    x = r * np.cos(theta)
    y = r * np.sin(theta)
    
    # 模拟高度变化和一些噪声
    z = np.random.normal(-0.2, 0.5, num_points)
    
    # 添加一些特异点和噪声
    outlier_mask = np.random.random(num_points) < 0.05  # 5% 的离群点
    x[outlier_mask] += np.random.normal(0, 2.0, np.sum(outlier_mask))
    y[outlier_mask] += np.random.normal(0, 2.0, np.sum(outlier_mask))
    z[outlier_mask] += np.random.normal(0, 2.0, np.sum(outlier_mask))
    
    # 组合为点云
    point_cloud = np.column_stack([x, y, z])
    
    return point_cloud


def main():
    parser = argparse.ArgumentParser(description="实时点云Transformer去噪工具")
    
    # 基本参数
    parser.add_argument('--model_path', type=str, required=True,
                        help='预训练模型路径')
    parser.add_argument('--device', type=str, default='auto',
                        choices=['auto', 'cuda', 'cpu'],
                        help='运行设备 (auto/cuda/cpu)')
    
    # 实时参数
    parser.add_argument('--simulation', action='store_true',
                        help='使用模拟的LiDAR数据')
    parser.add_argument('--lidar_device', type=str, default=None,
                        help='LiDAR设备路径')
    parser.add_argument('--buffer_size', type=int, default=5,
                        help='处理缓冲区大小')
    parser.add_argument('--fps', type=float, default=10.0,
                        help='模拟数据的帧率')
    
    # 可视化参数
    parser.add_argument('--visualization', action='store_true',
                        help='启用实时可视化')
    parser.add_argument('--save_frames', action='store_true',
                        help='保存处理后的帧')
    parser.add_argument('--output_dir', type=str, default='./results/realtime',
                        help='保存帧的输出目录')
    
    args = parser.parse_args()
    
    # 创建输出目录
    if args.save_frames:
        os.makedirs(args.output_dir, exist_ok=True)
    
    # 创建实时去噪器
    realtime_denoiser = RealTimePointCloudDenoiser(
        model_path=args.model_path,
        device=args.device,
        buffer_size=args.buffer_size,
        visualization=args.visualization
    )
    
    # 定义退出信号处理
    def signal_handler(sig, frame):
        print("\n检测到Ctrl+C，停止处理...")
        realtime_denoiser.stop()
        sys.exit(0)
    
    # 设置信号处理器
    signal.signal(signal.SIGINT, signal_handler)
    
    # 启动去噪处理
    realtime_denoiser.start()
    
    # 如果是模拟模式
    if args.simulation:
        try:
            frame_count = 0
            
            print("使用模拟的LiDAR数据进行实时去噪")
            print("按 'q' 键停止程序")
            
            while realtime_denoiser.running:
                # 检查是否按下q键
                if keyboard.is_pressed('q'):
                    print("\n检测到退出键，停止处理...")
                    break
                
                # 生成模拟数据
                point_cloud = simulate_lidar_data(num_points=np.random.randint(800, 1200))
                
                # 定义保存帧的回调函数
                def frame_callback(denoised_points):
                    nonlocal frame_count
                    if args.save_frames:
                        output_file = os.path.join(args.output_dir, f"frame_{frame_count:06d}.npy")
                        np.save(output_file, denoised_points)
                
                # 添加到处理队列
                realtime_denoiser.add_point_cloud(point_cloud, callback=frame_callback)
                
                # 帧计数器增加
                frame_count += 1
                
                # 按照指定帧率延时
                time.sleep(1.0 / args.fps)
        
        finally:
            # 确保在退出时停止处理
            realtime_denoiser.stop()
    
    # 使用真实LiDAR设备
    elif args.lidar_device:
        try:
            # 这里应该添加与特定LiDAR设备通信的代码
            # 由于不同LiDAR设备的API各不相同，这里只提供一个占位符
            print(f"连接到LiDAR设备 {args.lidar_device}")
            print("此功能需要根据特定设备实现")
            print("请参考您的LiDAR设备文档进行开发")
            
            # 假设有一个lidar_device对象，可以读取点云数据
            # while realtime_denoiser.running:
            #     point_cloud = lidar_device.read_frame()
            #     realtime_denoiser.add_point_cloud(point_cloud)
            #     time.sleep(0.01)  # 小延时避免CPU过载
            
            # 由于没有实际设备，这里简单等待用户退出
            print("按 'q' 键停止程序")
            while realtime_denoiser.running:
                if keyboard.is_pressed('q'):
                    print("\n检测到退出键，停止处理...")
                    break
                time.sleep(0.1)
        
        finally:
            # 确保在退出时停止处理
            realtime_denoiser.stop()
    
    # 如果没有指定数据源
    else:
        print("错误: 必须指定 --simulation 或 --lidar_device")
        realtime_denoiser.stop()
        sys.exit(1)
    
    print("处理完成!")


if __name__ == "__main__":
    main() 